﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace CRFEncoder
{
    public class LBFGS
    {
        float[] diag;
        BigArray<float> w;
        Mcsrch mcsrch_;
        long nfev, point, npt, iter, info, ispt, iypt;
        int iflag_;
        float stp;
        public int zeroone;
        public int err;
        public float obj;
        public BigArray<float> expected;

        public LBFGS()
        {
            iflag_ = 0; nfev = 0;
            point = 0; npt = 0; iter = 0; info = 0;
            ispt = 0; iypt = 0;
            stp = 0.0f;
            mcsrch_ = new Mcsrch();
        }


        public int optimize(long size, ref BigArray<float> x, float C)
        {
            const long msize = 5;
            if (w == null || w.Length == 0)
            {
                iflag_ = 0;
                w = new BigArray<float>(size * (2 * msize + 1) + 2 * msize, 1);
                diag = new float[size + 1];
            }
            else if (diag.Length != size + 1)
            {
                Console.WriteLine("size of array is different!");
                return -1;
            }

            lbfgs_optimize(size, msize, ref x, C);
            if (iflag_ < 0)
            {
                Console.WriteLine("routine stops with unexpected error");
                return -1;
            }

            return iflag_;
        }

        void lbfgs_optimize(long size,
                           long msize,
                           ref BigArray<float> x,
                           float C)
        {
            float yy = 0.0f;
            float ys = 0.0f;
            long bound = 0;
            long cp = 0;
            bool bExit = false;

            // initialization
            if (iflag_ == 0)
            {
                point = 0;
                ispt = size + (msize << 1);
                iypt = ispt + size * msize;
                Parallel.For(1, size + 1, i =>
                {
                    diag[i] = 1.0f;
                    w[ispt + i] = -expected[i];
                    w[i] = expected[i];
                });
                stp = 1.0f / (float)Math.Sqrt(Utils.ddot_(size, expected, 1, expected, 1));

                ++iter;
                info = 0;
                nfev = 0;
            }

            // MAIN ITERATION LOOP
            bExit = LineSearchAndUpdateStepGradient(size, msize, ref x);
            while (bExit == false)
            {
                ++iter;
                info = 0;
                if (iter > size)
                {
                    bound = size;
                }

                // COMPUTE -H*G USING THE FORMULA GIVEN IN: Nocedal, J. 1980,
                // "Updating quasi-Newton matrices with limited storage",
                // Mathematics of Computation, Vol.24, No.151, pp. 773-782.
                ys = Utils.ddot_(size, w, iypt + npt + 1, w, ispt + npt + 1);
                yy = Utils.ddot_(size, w, iypt + npt + 1, w, iypt + npt + 1);
                Parallel.For(1, size + 1, i =>
                {
                    diag[i] = ys / yy;
                    w[i] = -expected[i];
                });

                cp = point;
                if (point == 0)
                {
                    cp = msize;
                }

                w[size + cp] = 1.0f / ys;
                bound = Math.Min(iter - 1, msize);
                cp = point;
                for (int i = 1; i <= bound; ++i)
                {
                    --cp;
                    if (cp == -1) cp = msize - 1;
                    float sq = Utils.ddot_(size, w, ispt + cp * size + 1, w, 1);
                    long inmc = size + msize + cp + 1;
                    long iycn = iypt + cp * size;
                    w[inmc] = w[size + cp + 1] * sq;
                    float d = -w[inmc];

                    Parallel.For(1, size + 1, j =>
                        {
                            w[j] += d * w[iycn + j];
                        });
                }

                Parallel.For(1, size + 1, i =>
                {
                    w[i] = diag[i] * w[i];
                });

                for (int i = 1; i <= bound; ++i)
                {
                    float yr = Utils.ddot_(size, w, iypt + cp * size + 1, w, 1);
                    float beta = w[size + cp + 1] * yr;
                    long inmc = size + msize + cp + 1;
                    beta = w[inmc] - beta;
                    long iscn = ispt + cp * size;

                    Parallel.For(1, size + 1, j =>
                        {
                            w[j] += beta * w[iscn + j];
                        });

                    ++cp;
                    if (cp == msize)
                    {
                        cp = 0;
                    }
                }

                // STORE THE NEW SEARCH DIRECTION
                long offset = ispt + point * size;
                Parallel.For(1, size + 1, i =>
                {
                    w[offset + i] = w[i];
                    w[i] = expected[i];
                });

                stp = 1.0f;
                nfev = 0;
                bExit = LineSearchAndUpdateStepGradient(size, msize, ref x);
            }
        }

        private bool LineSearchAndUpdateStepGradient(long size, long msize, ref BigArray<float> x)
        {
            bool bExit = false;
            mcsrch_.mcsrch(size, x, obj, ref expected, w, ispt + point * size,
                            ref stp, ref info, ref nfev, diag);
            if (info == -1)
            {
                iflag_ = 1;  // next value
                bExit = true;
            }
            else if (info != 1)
            {
                //MCSRCH error, please see error code in info
                iflag_ = -1;
                bExit = true;
            }
            else
            {
                // COMPUTE THE NEW STEP AND GRADIENT CHANGE
                npt = point * size;
                Parallel.For(1, size + 1, i =>
                {
                    w[ispt + npt + i] = stp * w[ispt + npt + i];
                    w[iypt + npt + i] = expected[i] - w[i];
                });

                ++point;
                if (point == msize)
                {
                    point = 0;
                }

                double gnorm = Math.Sqrt(Utils.ddot_(size, expected, 1, expected, 1));
                double xnorm = Math.Max(1.0, Math.Sqrt(Utils.ddot_(size, x, 1, x, 1)));
                if (gnorm / xnorm <= Utils.eps)
                {
                    iflag_ = 0;  // OK terminated
                    bExit = true;
                }
            }

            return bExit;
        }
    }
}
